package com.demo.newLSTM;

import org.deeplearning4j.datasets.iterator.impl.ListDataSetIterator;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.GradientNormalization;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.layers.*;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.nd4j.linalg.activations.Activation;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.factory.Nd4j;
import org.nd4j.linalg.indexing.NDArrayIndex;
import org.nd4j.linalg.learning.config.Adam;
import org.nd4j.linalg.lossfunctions.LossFunctions;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

public class LSTMPredictionDemo1 {
    public static void main(String[] args) throws IOException {
        // 假设已经从数据库或其他数据源中获取到了负荷数据、温度数据和湿度数据等
        double[] load = {10.0, 12.0, 15.0, 20.0, 18.0, 14.0, 12.0, 10.0, 9.0, 8.0, 9.0, 11.0, 13.0, 15.0, 18.0, 20.0};
        double[] temperature = {20.0, 22.0, 25.0, 28.0, 30.0, 28.0, 25.0, 22.0, 20.0, 18.0, 16.0, 18.0, 20.0, 22.0, 25.0, 28.0};
        double[] humidity = {50.0, 55.0, 60.0, 65.0, 70.0, 75.0, 80.0, 85.0, 90.0, 85.0, 80.0, 75.0, 70.0, 65.0, 60.0, 55.0};

        // 准备输入数据
        int numExamples = load.length - 4;
        int numInputs = 4;
        INDArray input = Nd4j.zeros(numExamples, numInputs);
        //List<DataSet> dataList = new ArrayList<>();

        for (int i = 0; i < numExamples; i++) {
            input.putScalar(new int[] {i, 0}, load[i]);
            input.putScalar(new int[] {i, 1}, temperature[i]);
            input.putScalar(new int[] {i, 2}, humidity[i]);
            input.putScalar(new int[] {i, 3}, i % 24 + 1); // 添加时间信息，以便模型学习每天的周期性

            //dataList.add(new DataSet(Nd4j.create(input, new int[]{1, numInputs, 1}), Nd4j.create(output, new int[]{1, numOutputs, 1})));
        }



        // 准备标签数据
        int numOutputs = 1;
        INDArray labels = Nd4j.zeros(numExamples, numOutputs);
        for (int i = 4; i < load.length; i++) {
            labels.putScalar(new int[] {i - 4, 0}, load[i]);
        }

        // 构建模型
        int lstmLayerSize = 10;
        int numHiddenNodes = 10;
        int numTimeSteps = (int) input.size(0);
        LSTM lstmBuilder = new LSTM.Builder()
                .nIn(numInputs)
                .nOut(lstmLayerSize)
                .activation(Activation.TANH)
                .gateActivationFunction(Activation.SIGMOID)
                .dropOut(0.5)
                .build();
        MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder()
                .seed(123)
                .weightInit(WeightInit.XAVIER)
                .updater(new Adam(1e-3))
                .list()
                .layer(0, lstmBuilder)
                .layer(new RnnOutputLayer.Builder(LossFunctions.LossFunction.MSE)
                        .activation(Activation.IDENTITY)
                        .nIn(lstmLayerSize)
                        .nOut(numOutputs)
                        .build())
                .build();
        MultiLayerNetwork net = new MultiLayerNetwork(conf);
        net.init();

        // 训练模型
        int numEpochs = 100;
        for (int i = 0; i < numEpochs; i++) {
            INDArray reshape = input.reshape(numExamples, numTimeSteps / numExamples, numInputs);
            net.fit(reshape, labels);
        }

        // 预测未来负荷
        INDArray futureInput = Nd4j.zeros(1, numInputs);
        futureInput.putScalar(new int[] {0, 0}, load[load.length - 1]);
        futureInput.putScalar(new int[] {0, 1}, temperature[temperature.length - 1]);
        futureInput.putScalar(new int[] {0, 2}, humidity[humidity.length - 1]);
        futureInput.putScalar(new int[] {0, 3}, (load.length % 24) + 1);
        List<Double> predictions = new ArrayList<>();
        for (int i = 0; i < 24; i++) {
            INDArray output = net.output(futureInput);
            //double prediction = output.getDouble(NDArrayIndex.point(0), NDArrayIndex.point(0));
            double prediction = output.getDouble(0, 0, i);

            predictions.add(prediction);
            futureInput.putScalar(new int[] {0, 0}, prediction);
            futureInput.putScalar(new int[] {0, 3}, (futureInput.getDouble(0,3) + 1) % 24 + 1);
        }

        // 输出预测结果
        System.out.println("预测结果：");
        for (int i = 1; i <= predictions.size(); i++) {
            System.out.println("第" + i + "小时：" + predictions.get(i - 1));
        }
    }
}
